import streamlit as st
from PIL import Image
import base64
from openai import OpenAI

openai_api_key = "EMPTY"
openai_api_base = "https://u359243-a7bc-31824cd4.westx.seetacloud.com:8443/v1"  # 替换成自己的地址

model_name = "qwen-vl"  # 模型名称

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)

st.set_page_config(
    page_title="QwenVL识别",
    page_icon=" ",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Title and description in main area
st.title("QwenVL识别")

# Add clear button to top right
col1, col2 = st.columns([6, 1])
with col2:
    if st.button("清除  ️"):
        if 'ocr_result' in st.session_state:
            del st.session_state['ocr_result']
        st.rerun()

st.markdown('<p style="margin-top: -20px;">QWenVL-7B图像处理</p>',
            unsafe_allow_html=True)
st.markdown("---")

# Move upload controls to sidebar
with st.sidebar:
    st.header("上传图片")
    uploaded_file = st.file_uploader("选择图片...", type=['png', 'jpg', 'jpeg'])

    if uploaded_file is not None:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="上传图像")

        if st.button("运行", type="primary"):
            with st.spinner("处理中..."):
                try:
                    image_result = base64.b64encode(uploaded_file.getvalue()).decode('utf-8')
                    chat_completion_from_base64 = client.chat.completions.create(
                        messages=[{
                            "role":
                                "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": """：
你是Qwen-VL模型助手，接收图片后请先完成图像类别识别（文本类/表格类/流程图/其他图像），并按以下规则处理：
- **文本类图像**：需精准识别所有可见文本，完整保留原始排版结构（如段落缩进、对齐方式、字体层级等格式），仅输出纯转录内容，不添加任何主观评论或解析。
- **表格类图像**：需识别表格行列结构（含合并单元格、跨列等复杂格式），提取单元格内文字信息，以标准HTML表格语法输出，确保行列对应关系与原图一致。
- **流程图**：需解析图形元素（矩形/菱形/箭头等）逻辑关系，转化为Mermaid语法的Markdown文本，包含节点定义、连接关系及箭头流向说明，保留原图流程逻辑。
- **其他图像**：需从客观视角描述图像内容，包含主体元素、色彩基调、场景特征、关键细节（如尺寸比例、位置关系等），若存在可识别文本需同步标注。
"""
                                },
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:image/jpeg;base64,{image_result}"
                                    },
                                },
                            ],
                        }],
                        model=model_name,
                        max_completion_tokens=1024,
                    )

                    result = chat_completion_from_base64.choices[0].message.content
                    st.session_state['ocr_result'] = result
                except Exception as e:
                    st.error(f"Error processing image: {str(e)}")

# Main content area for results
if 'ocr_result' in st.session_state:
    st.markdown(st.session_state['ocr_result'])
else:
    st.info("上传图像并点击'运行'按钮进行识别.")

# Footer
st.markdown("---")
st.markdown("by-Qwen2.5-Vl-7B 模型")